import datetime
import pandas as pd
import tushare as ts
import matplotlib.pyplot as plt
from matplotlib.widgets import MultiCursor
from pandas.plotting import register_matplotlib_converters

import QUANTAXIS as QA


register_matplotlib_converters()


class BottomBox(object):
    def __init__(self, data, high_point=None, low_point=None, init_ratio=0.4):
        self.data = pd.DataFrame()
        self.tail_data = pd.DataFrame()
        self.all_data = data.loc[high_point['index']:]
        self.filter_ratio = init_ratio
        self.real_ratio = init_ratio
        self.highest = high_point
        self.lowest = low_point
        self.high = (high_point['value'] - low_point['value']) * self.filter_ratio + low_point['value']
        # print('底部高点设置为：%.2f' % self.high)

    @property
    def length(self):
        return len(self.data)

    @property
    def tail_length(self):
        return len(self.tail_data)

    @property
    def tail_high_date(self):
        if self.length > 0:
            return self.tail_data['close'].idxmax()
        else:
            return 'empty'

    @property
    def start_date(self):
        if self.length > 0:
            return self.data.index[0]
        else:
            return 'empty'

    @property
    def end_date(self):
        if self.length > 0:
            return self.data.index[-1]
        else:
            return 'empty'

    @property
    def not_through(self):
        if str(self.data.index[-1]) == str(self.all_data.index[-1]):
            return True
        else:
            return False

    def filter(self):
        find_lowest = False
        for index, row in self.all_data.iterrows():
            # 当高点比框体上沿小时，加入数据
            if row['high'] <= self.high:  # 小于等于的意思是包含最高点数据，如果没有等于，会清空之前的数据
                self.data = self.data.append(row)
                if index == self.lowest['index']:
                    find_lowest = True
            elif find_lowest:  # 当高点超出框体上沿时，如果已包含最低点，则结束框体
                self.tail_data = self.all_data.loc[index:]
                break
            else:
                self.data = pd.DataFrame()

    def compress_height(self, invalid_count=3):
        '''
        压缩高度，找到近似中间最高点
        :param invalid_count: 两头无效数据的数量
        :return:
        '''
        if self.length < invalid_count * 2 + 1:
            return
        df = self.data.iloc[invalid_count: -invalid_count]
        high = df['high'].max()
        high_index = df['high'].idxmax()
        # print('截取前长度：%d---截取后长度：%d' % (len(self.data), len(df)))
        # print('高点位于%s' % str(high_index))
        if high_index in df.iloc[:invalid_count].index or high_index in df.iloc[-invalid_count:].index:
            self.data = df
            self.compress_height(invalid_count=invalid_count)
        else:
            self.high = high
            self.data = pd.DataFrame()
            self.filter()

    def tail_through(self, invalid_count=2):
        if not self.not_through and self.tail_length > 5:
            self.tail_data = self.tail_data.iloc[invalid_count:]
            high = self.tail_data['close'].max()
            high_index = self.tail_data['close'].idxmax()
            low = self.tail_data['close'].min()
            low_index = self.tail_data['close'].idxmin()

            if low < self.high * 0.9 or high < self.high:
                return False
            if high_index > low_index:
                return True
            else:
                return False
        else:
            return False


def main(codes, start, end):
    # print(codes.head())
    data = QA.QA_fetch_stock_day_adv(codes['code'].to_list(), start, end)
    data = data.to_qfq()
    data = data.week
    df_result = pd.DataFrame(columns=['code', 'start', 'end', 'length', 'high'])
    for _, code in codes.iterrows():
        # print(code)
        # df = data.data.loc[(slice(None), code), :]
        df = data.data.xs(code['code'], level='code', axis=0, drop_level=True)
        box = get_box(df)

        if box.length > 50 and box.tail_through() \
                and box.tail_high_date.year == 2021 \
                and box.tail_high_date.month >= 4\
                and box.tail_high_date.day >= 1:
            print('code：{}---{}---起始时间{}---终止时间{}---长度{}---高值{}---尾部高值时间 {}'.format(
                code['code'], code['name'], box.start_date, box.end_date, box.length, box.high, box.tail_high_date))
            df_result = df_result.append({'code': code['code'],
                                          'start': box.start_date,
                                          'end': box.end_date,
                                          'length': box.length,
                                          'high': box.high}, ignore_index=True)

            figure = plt.figure()
            plt.title('{}---'.format(code['code']))
            axes1 = figure.add_subplot(411)
            axes2 = figure.add_subplot(412)
            axes3 = figure.add_subplot(413)
            axes4 = figure.add_subplot(414)
            axes1.plot(df.index, df['high'])
            axes2.plot(df.index, df['volume'])
            axes3.plot(df.index, df['close'])
            axes4.plot(df.index, df['high'])

            left, right = axes1.get_xlim()
            axes1.hlines(y=box.high, xmin=left, xmax=right, linestyles='dashed')
            multi = MultiCursor(figure.canvas, (axes1, axes2, axes3, axes4), color='r', lw=1)

            plt.show()
    df_result.to_csv('break_through.csv')

    # data.to_csv('002581.csv')
    # print(data.index)


def get_box(data: pd.DataFrame):
    max_point = {'index': data['high'].idxmax(), 'value': data['high'].max()}
    min_point = {'index': data['low'].idxmin(), 'value': data['low'].min()}
    box = BottomBox(data=data, high_point=max_point, low_point=min_point)
    box.filter()
    # ratio = 0.2
    # count = int(box.length * ratio)
    box.compress_height(invalid_count=3)
    return box


if __name__ == '__main__':
    ts_token = '17056d23a59ab71cb979c6a30185e092aba605c4544dac900a3eb7f8'
    ts.set_token(ts_token)
    pro = ts.pro_api()
    data = pro.stock_basic(exchange='', list_status='L', fields='symbol,name,area,industry,list_date')
    data = data[data.symbol.str.startswith('00') | data.symbol.str.startswith('60') | data.symbol.str.startswith('30')]
    # data = data[data.symbol.str.startswith('002068')]
    data = data[data.list_date < '20160101']
    all = data.rename(columns={'symbol': 'code'})
    codes = all[~all.name.str.contains('ST')]
    # codes = codes.iloc[0: 1]
    # hs300 = pd.read_csv('D:\\PythonPro\\QUANTAXIS\\AI\\hs300.csv')
    # hs300 = ts.get_hs300s()
    # codes = codes['code'].to_list()
    # codes = ['002027']
    # print(codes)
    start = '2016-01-01'
    end = '2021-04-09'
    main(codes, start, end)
